from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection



# 載入 "./logs/max_le_collection.npy" 和 "./logs/max_ler_collection.npy" 這兩個矩陣

max_le_collection = np.load("./logs/max_le_collection.npy")
max_ler_collection = np.load("./logs/max_ler_collection.npy")

print("shape of max_le_collection: ", max_le_collection.shape)
print("shape of max_ler_collection: ", max_ler_collection.shape)

# # 將 max_le_collection 和 max_ler_collection 中的所有記錄繪製出來
# fig = plt.figure()
# for i in range(max_le_collection.shape[0]):
#     # 用每一条线的 i 值添加标签
#     plt.plot(max_ler_collection[i], label=str(i)+'_ler')
#     plt.plot(max_le_collection[i], label=str(i)+'_le')
# plt.legend()
# plt.show()

# 繪製 max_le_collection 和 max_ler_collection 各自的包絡

max_ler_collection_upper_bound = np.max(max_ler_collection, axis=0)
max_le_collection_upper_bound = np.max(max_le_collection, axis=0)
max_ler_collection_lower_bound = np.min(max_ler_collection, axis=0)
max_le_collection_lower_bound = np.min(max_le_collection, axis=0)

epsilons = [(i+1)*0.01 for i in range(max_le_collection.shape[1])]

# 绘制包络图，使用与e_r64a.py SVG版本完全一致的颜色和字体方案
fig = plt.figure(figsize=(12, 8))

# 使用与e_r64a.py相同的颜色方案：蓝色为训练模型，橙色为随机模型
plt.fill_between(epsilons, max_ler_collection_lower_bound, max_ler_collection_upper_bound, color='orange', alpha=0.3, label='Random Model Range')
plt.fill_between(epsilons, max_le_collection_lower_bound, max_le_collection_upper_bound, color='blue', alpha=0.3, label='Trained Model Range')

# 绘制均值线，使用相同的颜色和线宽
plt.plot(epsilons, np.mean(max_ler_collection, axis=0), color='orange', linewidth=3, label='Random Model Mean')
plt.plot(epsilons, np.mean(max_le_collection, axis=0), color='blue', linewidth=3, label='Trained Model Mean')

# 在x=0.375位置绘制蓝色竖线，标注invariant neighborhood
plt.axvline(x=0.375, color='blue', linewidth=2, linestyle=':', alpha=0.8, label='Invariant Neighborhood')
plt.text(0.375, plt.ylim()[1]*0.8, 'Invariant\nNeighborhood', fontsize=12, ha='center', 
         bbox=dict(boxstyle='round,pad=0.3', facecolor='lightblue', alpha=0.7))

# 使用与e_r64a.py完全相同的字体方案
plt.xlabel('Epsilon', fontsize=14)
plt.ylabel('Primary Lyapunov Exponent', fontsize=14)
plt.title('Lyapunov Exponent vs Perturbation Magnitude\n(Trained vs Random Model)', fontsize=16)
plt.legend(fontsize=12)
plt.grid(True, alpha=0.3)

plt.tight_layout()

# 保存PNG和SVG版本
plt.savefig('./logs/lyapunov_envelope.png', dpi=300, bbox_inches='tight')

# 生成SVG版本
fig_svg = plt.figure(figsize=(12, 8))

# 重新绘制相同的图表用于SVG输出（完全相同的参数）
plt.fill_between(epsilons, max_ler_collection_lower_bound, max_ler_collection_upper_bound, color='orange', alpha=0.3, label='Random Model Range')
plt.fill_between(epsilons, max_le_collection_lower_bound, max_le_collection_upper_bound, color='blue', alpha=0.3, label='Trained Model Range')

plt.plot(epsilons, np.mean(max_ler_collection, axis=0), color='orange', linewidth=3, label='Random Model Mean')
plt.plot(epsilons, np.mean(max_le_collection, axis=0), color='blue', linewidth=3, label='Trained Model Mean')

# 在SVG版本中也添加相同的invariant neighborhood标注
plt.axvline(x=0.375, color='blue', linewidth=2, linestyle=':', alpha=0.8, label='Invariant Neighborhood')
plt.text(0.375, plt.ylim()[1]*0.8, 'Invariant\nNeighborhood', fontsize=12, ha='center', 
         bbox=dict(boxstyle='round,pad=0.3', facecolor='lightblue', alpha=0.7))

plt.xlabel('Epsilon', fontsize=14)
plt.ylabel('Primary Lyapunov Exponent', fontsize=14)
plt.title('Lyapunov Exponent vs Perturbation Magnitude\n(Trained vs Random Model)', fontsize=16)
plt.legend(fontsize=12)
plt.grid(True, alpha=0.3)

plt.tight_layout()
fig_svg.savefig('./logs/lyapunov_envelope.svg', format='svg', bbox_inches='tight')
plt.close(fig_svg)

plt.show()

print("Results saved:")
print("- PNG envelope plot saved as ./logs/lyapunov_envelope.png")
print("- SVG envelope plot saved as ./logs/lyapunov_envelope.svg")
